{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec83ccb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(os.path.abspath(os.path.pardir))\n",
    "from argparse import ArgumentParser\n",
    "\n",
    "import torch\n",
    "from voltaml.compile import VoltaGPUCompiler\n",
    "from voltaml.inference import gpu_performance\n",
    "import torchvision"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff3892df",
   "metadata": {},
   "source": [
    "### Load Model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ccc4c2-dd5c-4f8a-8eec-e60ce8392eb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_mobilenet_v3_large', pretrained=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91758d3d",
   "metadata": {},
   "source": [
    "## Set parameters for FP16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad6b0369",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape = (1,3,224,224)\n",
    "precision = 'fp16'\n",
    "compiled_model_dir = 'deeplabv3_mobile.engine' ## Set Model dir\n",
    "throughput_batch_size = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58f4f365",
   "metadata": {},
   "source": [
    "### Compile Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d078430b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "compiler = VoltaGPUCompiler(\n",
    "    model=model,\n",
    "    output_dir=compiled_model_dir,\n",
    "    input_shape=input_shape,\n",
    "    precision=precision,\n",
    "    simplify=True\n",
    ")\n",
    "\n",
    "compiled_model = compiler.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4fdbf14",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_performance(compiled_model_dir, model, input_shape=input_shape, throughput_batch_size=throughput_batch_size, is_yolo=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3341b48d",
   "metadata": {},
   "source": [
    "### Set parameters for INT8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a2a0731",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape = (1,3,224,224)\n",
    "precision = 'int8'\n",
    "compiled_model_dir = 'deeplabv3_mobile_int8.engine' ## Compiled model directory\n",
    "throughput_batch_size = 1\n",
    "calib_input = '/workspace/val2017/' ## Calib input images path\n",
    "calib_cache = '' ## Cache name\n",
    "calib_num_images=25000\n",
    "calib_batch_size=8\n",
    "calib_preprocessor='V1MS'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e966cfb9",
   "metadata": {},
   "source": [
    "### Compile Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b0a64c5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "compiler = VoltaGPUCompiler(\n",
    "    model=model,\n",
    "    output_dir=compiled_model_dir,\n",
    "    input_shape=input_shape,\n",
    "    precision=precision,\n",
    "    calib_input=calib_input,\n",
    "    calib_cache=calib_cache,\n",
    "    calib_num_images=calib_num_images,\n",
    "    calib_batch_size=calib_batch_size,\n",
    "    calib_preprocessor=calib_preprocessor\n",
    ")\n",
    "\n",
    "compiled_model = compiler.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a22fb82",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_performance(compiled_model_dir, model, input_shape=input_shape, throughput_batch_size=throughput_batch_size, is_yolo=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea3116e0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd8a5daa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1422ab34",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
